package _18_剑指OfferII;

import java.util.*;

public class _061_剑指OfferII和最小的k个数对 {
    
    // 初始化最小推
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        PriorityQueue<int[]> minQueue = new PriorityQueue<>((a1, a2) -> nums1[a1[0]] + nums2[a1[1]] - nums1[a2[0]] - nums2[a2[1]]);
        // 初始化最小推
        int len1 = Math.min(nums1.length, k);
        for (int i = 0; i < len1; ++i) {
            minQueue.add(new int[]{i, 0});
        }

        ArrayList<List<Integer>> results = new ArrayList<>();
        while (k-- > 0 && !minQueue.isEmpty()) {
            int[] poll = minQueue.poll();
            results.add(Arrays.asList(nums1[poll[0]], nums2[poll[1]]));

            if (poll[1] < nums2.length - 1) {
                minQueue.add(new int[]{poll[0], poll[1] + 1});
            }
        }
        return results;
    }

    // Topk 问题
    public List<List<Integer>> kSmallestPairs1(int[] nums1, int[] nums2, int k) {
        PriorityQueue<int[]> maxQueue = new PriorityQueue<>((a1, a2) -> a2[0] + a2[1] - a1[0] - a1[1]);

        // 将问题装化为k^2logk的问题
        int len1 = Math.min(nums1.length, k);
        int len2 = Math.min(nums2.length, k);

        for (int i = 0; i < len1; ++i) {
            for (int j = 0; j < len2; ++j) {
                // 不应该直接插入，直接插入存在多余的上滤
                if (maxQueue.size() >= k) {
                    int[] root = maxQueue.peek();
                    if (root[0] + root[1] > nums1[i] + nums2[j]) {
                        maxQueue.poll();
                        maxQueue.offer(new int[] {nums1[i], nums2[j]});
                    }
                } else {
                    maxQueue.offer(new int[] {nums1[i], nums2[j]});
                }
            }
        }
        List<List<Integer>> result = new ArrayList<>();
        while (!maxQueue.isEmpty()) {
            int[] poll = maxQueue.poll();
            result.add( Arrays.asList(poll[0], poll[1]));
        }
        return result;
    }

}
